#pragma once

#include "task/Task.hpp"

#include <queue>
#include <mutex>
#include <memory>
#include <atomic>
#include <condition_variable>

namespace threadpool {

    class TaskQueue {

    public:

        // 声明任务比较器类型
        using TaskComparator = std::function<bool(const std::shared_ptr<Task> &, const std::shared_ptr<Task> &)>;

        /**
          * @brief  任务队列有参构造函数
          * @note   None
          * @param  strategy 继承自TaskSchedulerStrategy类的子类实例化对象
          */
        template<typename CompareStrategy>
        explicit TaskQueue(const CompareStrategy &strategy) : queue_(strategy), stop_(false) {}

        /**
          * @brief  插入一个任务
          * @note   线程安全
          * @param  task 任务的智能指针
          */
        void push(std::shared_ptr<Task> task);

        /**
          * @brief  获取任务队列队头元素
          * @note   线程安全
          * @param  None
          * @retval 任务队列队头元素智能指针
          */
        std::shared_ptr<Task> top();

        /**
          * @brief  弹出任务队列一个对头元素
          * @note   线程安全，如果任务队列内没有元素则会阻塞，直至有新的元素才会继续弹出
          * @param  None
          * @retval 任务队列队头元素智能指针
          */
        std::shared_ptr<Task> pop();

        /**
          * @brief  检查任务队列是否为空
          * @note   线程安全
          * @param  None
          * @retval 任务队列是否为空
          */
        bool empty();

        /**
          * @brief  获取任务队列队长
          * @note   线程安全
          * @param  None
          * @retval 任务队列队长
          */
        size_t size();

        /**
          * @brief  关闭任务队列
          * @note   这将导致所有被阻塞在pop方法的线程直接返回nullptr，并且该线程池将被置为关闭状态无法再使用
          * @param  None
          */
        void stop();

    private:

        // 任务队列
        std::priority_queue<std::shared_ptr<Task>, std::vector<std::shared_ptr<Task>>, TaskComparator> queue_;

        // 停止标志
        std::atomic<bool> stop_;

        // 互斥锁
        std::mutex mutex_;

        // 条件变量
        std::condition_variable cond_;

    };

    void TaskQueue::push(std::shared_ptr<Task> task) {
        std::lock_guard<std::mutex> lock(mutex_);
        queue_.push(std::move(task));
        cond_.notify_one();
    }

    std::shared_ptr<Task> TaskQueue::top() {
        std::lock_guard<std::mutex> lock(mutex_);
        return queue_.top();
    }

    std::shared_ptr<Task> TaskQueue::pop() {
        std::unique_lock<std::mutex> lock(mutex_);
        cond_.wait(lock, [this]() { return !queue_.empty() || stop_.load(); });
        if (stop_.load() || queue_.empty()) {
            return nullptr;
        }
        auto task = queue_.top();
        queue_.pop();
        return task;
    }

    bool TaskQueue::empty() {
        std::lock_guard<std::mutex> lock(mutex_);
        return queue_.empty();
    }

    size_t TaskQueue::size() {
        std::lock_guard<std::mutex> lock(mutex_);
        return queue_.size();
    }

    void TaskQueue::stop() {
        std::lock_guard<std::mutex> lock(mutex_);
        stop_.store(true);
        cond_.notify_all();
    }

}
